#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
@Project : hello
@file : AsyncOracleExecutor.py
@Author : sheen
@time : 2025/5/14 16:48
@func : 批量oracle异步查询demo


"""
import asyncio
import oracledb
from typing import List, Dict, Optional
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class AsyncOracleExecutor:
    """简化版异步Oracle查询执行器"""

    def __init__(self, user: str, password: str, dsn: str, pool_size: int = 5):
        self.pool = None
        self.config = {
            "user": user,
            "password": password,
            "dsn": dsn,
            "min": 1,
            "max": pool_size,
            "increment": 1
        }

    async def initialize(self):
        """初始化连接池"""
        try:
            # 自动选择Thin/Thick模式
            try:
                oracledb.defaults.driver_name = "thin"
                logger.info("使用Thin模式连接Oracle")
            except:
                oracledb.init_oracle_client()
                logger.info("使用Thick模式连接Oracle")

            self.pool = await oracledb.create_pool_async(**self.config)
            logger.info(f"连接池初始化成功 (max={self.config['max']})")
        except Exception as e:
            logger.error(f"连接池初始化失败: {e}")
            raise

    async def execute_query(self, sql: str, params: Optional[Dict] = None) -> List[Dict]:
        """执行单个查询"""
        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                await cursor.execute(sql, params or {})

                # 获取结果
                columns = [col[0] for col in cursor.description]
                rows = await cursor.fetchall()

                # 转换为字典列表
                return [dict(zip(columns, row)) for row in rows]

    async def execute_parallel(self, queries: List[str], max_concurrent: int = 5) -> List[List[Dict]]:
        """并行执行多个查询"""
        semaphore = asyncio.Semaphore(max_concurrent)

        async def run_query(sql: str):
            async with semaphore:
                try:
                    result = await self.execute_query(sql)
                    logger.info(f"查询成功: {sql[:50]}...")
                    return result
                except Exception as e:
                    logger.error(f"查询失败: {sql[:50]}... - {str(e)}")
                    return []

        tasks = [asyncio.create_task(run_query(sql)) for sql in queries]
        return await asyncio.gather(*tasks)

    async def close(self):
        """关闭连接池"""
        if self.pool:
            await self.pool.close()
            logger.info("连接池已关闭")


async def main():
    # 配置数据库连接
    executor = AsyncOracleExecutor(
        user="your_username",
        password="your_password",
        dsn="host:1521/service_name",
        pool_size=5
    )

    try:
        await executor.initialize()

        # 示例查询
        queries = [
            "SELECT * FROM employees WHERE rownum <= 10",
            "SELECT department_id, COUNT(*) FROM employees GROUP BY department_id"
        ]

        results = await executor.execute_parallel(queries, max_concurrent=2)

        # 打印结果
        for i, (sql, data) in enumerate(zip(queries, results)):
            print(f"\n查询 {i + 1}: {sql[:50]}...")
            print(f"返回行数: {len(data)}")
            if data:
                print("示例数据:", data[0])

    finally:
        await executor.close()


if __name__ == "__main__":
    asyncio.run(main())